<!DOCTYPE html>
<html lang="en">

<head>
  <meta charset="UTF-8">
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
  <link rel="shortcut icon" href="data:;base64,iVBORw0KGgo=">
  <title>WebGPU Demo - step 16 - Instance</title>
  <style>
    html,
    body {
      margin: 0;
      overflow: hidden;
      height: 100%;
    }
  </style>
</head>

<body>
  <canvas id="canvas"></canvas>
  <script type="module">
    import { mat4, vec3, vec4, quat } from '../js/vendor/gl-matrix.js';
    import Camera from '../js/Camera.js';
    import { GUI } from '../js/vendor/dat.gui.js';
    import glslangModule from '../js/vendor/glslang.js';

    (async () => {
      if (!navigator.gpu) {
        const div = document.createElement('div');
        div.innerHTML = '<div>Your browser does not support WebGPU. Go to <a href="https://webgpu.io">webgpu.io</a> for more information.</div>';
        document.body.insertBefore(div, canvas);
        return;
      }

      // Initialization
      const adapter = await navigator.gpu.requestAdapter();
      const device = await adapter.requestDevice();
      const context = canvas.getContext('gpupresent');
      const format = await context.getSwapChainPreferredFormat(device);
      const swapChain = context.configureSwapChain({ device, format });

      // Create Camera Controller
      const camera = new Camera();
      camera.attach(canvas);

      const glslang = await glslangModule();

      const vertexShader = `
        layout(set = 0, binding = 0) uniform UniformScene {
          mat4 u_viewMatrix;
          mat4 u_projectionMatrix;
        };

        layout(location = 0) in vec3 a_position;
        layout(location = 1) in vec3 a_normal;
        layout(location = 2) in vec2 a_uv;
        // layout(location = 3) in mat4 a_transform;
        layout(location = 3) in vec4 a_transform0;
        layout(location = 4) in vec4 a_transform1;
        layout(location = 5) in vec4 a_transform2;
        layout(location = 6) in vec4 a_transform3;
        layout(location = 7) in vec4 a_color;

        layout(location = 0) out vec3 v_normal;
        layout(location = 1) out vec2 v_uv;
        layout(location = 2) out vec4 v_color;

        void main () {
          // gl_Position = u_projectionMatrix * u_viewMatrix * a_transform * vec4(a_position, 1.0);
          gl_Position = u_projectionMatrix * u_viewMatrix * mat4(a_transform0, a_transform1, a_transform2, a_transform3) * vec4(a_position, 1.0);
          v_normal = a_normal;
          v_uv = a_uv;
          v_color = a_color;
        }
      `;
      const fragmentShader = `
        #ifdef DIFFUSE_TEXTURE
          layout(set = 1, binding = 0) uniform sampler u_diffuseSampler;
          layout(set = 1, binding = 1) uniform texture2D u_diffuseTexture;
        #endif

        layout(location = 0) in vec3 v_normal;
        layout(location = 1) in vec2 v_uv;
        layout(location = 2) in vec4 v_color;

        layout(location = 0) out vec4 fragColor;

        void main () {
          fragColor = v_color;
          #ifdef DIFFUSE_TEXTURE
            vec4 diffuseTextureColor = texture(sampler2D(u_diffuseTexture, u_diffuseSampler), v_uv);
            fragColor *= diffuseTextureColor;
          #endif
        }
      `;

      // Create Render Pipeline
      const sampleCount = 4;
      function createPipeline(defines = []) {
        function getShaderWithDefines(shaderSource) {
          const defineStr = defines.map((define) => `#define ${define}`).join('\n');
          return `#version 450\n${defineStr}\n${shaderSource}`;
        }
        return device.createRenderPipeline({
          vertexStage: {
            module: device.createShaderModule({
              code: glslang.compileGLSL(getShaderWithDefines(vertexShader), 'vertex'),
            }),
            entryPoint: 'main',
          },
          fragmentStage: {
            module: device.createShaderModule({
              code: glslang.compileGLSL(getShaderWithDefines(fragmentShader), 'fragment'),
            }),
            entryPoint: 'main',
          },
          primitiveTopology: 'triangle-list',
          vertexState: {
            indexFormat: 'uint32',
            vertexBuffers: [
              {
                arrayStride: 4 * (3 + 3 + 2),
                attributes: [
                  {
                    format: 'float3',
                    offset: 0,
                    shaderLocation: 0,
                  },
                  {
                    format: 'float3',
                    offset: 4 * 3,
                    shaderLocation: 1,
                  },
                  {
                    format: 'float2',
                    offset: 4 * (3 + 3),
                    shaderLocation: 2,
                  },
                ],
              },
              {
                arrayStride: 4 * (4 * 4 + 4),
                stepMode: 'instance',
                attributes: [
                  {
                    format: 'float4',
                    offset: 0,
                    shaderLocation: 3,
                  },
                  {
                    format: 'float4',
                    offset: 4 * 4 * 1,
                    shaderLocation: 4,
                  },
                  {
                    format: 'float4',
                    offset: 4 * 4 * 2,
                    shaderLocation: 5,
                  },
                  {
                    format: 'float4',
                    offset: 4 * 4 * 3,
                    shaderLocation: 6,
                  },
                  {
                    format: 'float4',
                    offset: 4 * 4 * 4,
                    shaderLocation: 7,
                  },
                ],
              },
            ],
          },
          colorStates: [{ format }],
          depthStencilState: {
            format: 'depth24plus-stencil8',
            depthWriteEnabled: true,
            depthCompare: 'less-equal',
          },
          sampleCount,
        });
      }
      const pipelineWithTexture = createPipeline(['DIFFUSE_TEXTURE']);
      const pipelineWithoutTexture = createPipeline();

      const passDescriptor = {
        colorAttachments: [{
          attachment: undefined,
          loadValue: [0, 0, 0, 1],
        }],
        depthStencilAttachment: {
          attachment: undefined,
          depthLoadValue: 1,
          depthStoreOp: 'store',
          stencilLoadValue: 0,
          stencilStoreOp: 'store',
        },
      };

      // Create Output Texture
      let outputTexture;
      let outputTextureView;
      function createOutputTexture() {
        if (outputTexture) {
          outputTexture.destroy();
        }
        outputTexture = device.createTexture({
          size: {
            width: canvas.width,
            height: canvas.height,
            depth: 1,
          },
          sampleCount,
          format,
          usage: GPUTextureUsage.OUTPUT_ATTACHMENT,
        });
        outputTextureView = outputTexture.createView();
      }

      // Create Depth Texture
      let depthTexture;
      function createDepthTexture() {
        if (depthTexture) {
          depthTexture.destroy();
        }
        depthTexture = device.createTexture({
          size: {
            width: canvas.width,
            height: canvas.height,
            depth: 1,
          },
          sampleCount,
          format: 'depth24plus-stencil8',
          usage: GPUTextureUsage.OUTPUT_ATTACHMENT
        });
        passDescriptor.depthStencilAttachment.attachment = depthTexture.createView();
      }

      function resize() {
        const width = window.innerWidth;
        const height = window.innerHeight;
        const { devicePixelRatio } = window;
        canvas.width = width * devicePixelRatio;
        canvas.height = height * devicePixelRatio;
        canvas.style.width = `${width}px`;
        canvas.style.height = `${height}px`;

        createOutputTexture();
        createDepthTexture();
      }

      window.onresize = resize;
      resize();

      // Change background color
      const gui = new GUI();
      const config = {
        backgroundColor: [0, 0, 0],
        baseColor: [255, 255, 255],
      };
      gui.addColor(config, 'backgroundColor').onChange((value) => {
        passDescriptor.colorAttachments[0].loadValue[0] = value[0] / 255;
        passDescriptor.colorAttachments[0].loadValue[1] = value[1] / 255;
        passDescriptor.colorAttachments[0].loadValue[2] = value[2] / 255;
      });

      // Create Image Texture
      const image = new Image();
      image.src = '../asset/images/uv-grid.jpg';
      await image.decode();
      const imageBitmap = await createImageBitmap(image);
      const imageSize = {
        width: imageBitmap.width,
        height: imageBitmap.height,
        depth: 1,
      };
      const imageTexture = device.createTexture({
        size: imageSize,
        format: 'rgba8unorm',
        usage: GPUTextureUsage.COPY_DST | GPUTextureUsage.SAMPLED,
      });
      device.defaultQueue.copyImageBitmapToTexture(
        {
          imageBitmap,
        },
        {
          texture: imageTexture,
        },
        imageSize,
      );

      // Create Scene Buffer
      const sceneBuffer = device.createBuffer({
        size: 4 * (16 + 16),
        usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
      });

      // Create Vertex Buffer
      function createGeometryBuffer(vertexArray, indexArray) {
        const vertexBuffer = device.createBuffer({
          size: 4 * vertexArray.length,
          usage: GPUBufferUsage.VERTEX,
          mappedAtCreation: true,
        });
        const vertexArrayBuffer = vertexBuffer.getMappedRange();
        new Float32Array(vertexArrayBuffer).set(vertexArray);
        vertexBuffer.unmap();
        // Create Index Buffer
        const indexBuffer = device.createBuffer({
          size: 4 * indexArray.length,
          usage: GPUBufferUsage.INDEX,
          mappedAtCreation: true,
        });
        const indexArrayBuffer = indexBuffer.getMappedRange();
        new Uint32Array(indexArrayBuffer).set(indexArray);
        indexBuffer.unmap();
        return {
          vertexBuffer,
          indexBuffer,
          count: indexArray.length,
        };
      }
      const planeBuffers = createGeometryBuffer(
        [
          -1, -1, 0, 0, 0, 1, 0, 1,
          1, -1, 0, 0, 0, 1, 1, 1,
          -1, 1, 0, 0, 0, 1, 0, 0,
          1, 1, 0, 0, 0, 1, 1, 0,
        ],
        [
          0, 1, 2,
          1, 3, 2,
        ],
      );
      const triangleBuffers = createGeometryBuffer(
        [
          -1, -1, 0, 0, 0, 1, 0.0, 1.0,
          1, -1, 0, 0, 0, 1, 1.0, 1.0,
          0, 1, 0, 0, 0, 1, 0.5, 0.0,
        ],
        [
          0, 1, 2,
        ],
      );

      // Create Objects
      const objects = [
        {
          instances: [
            {
              rotation: quat.create(),
              translation: vec3.fromValues(-0.5, 0, -.5),
              scale: vec3.fromValues(0.3, 0.3, 0.3),
              color: vec4.fromValues(1, 0, 0, 1),
            },
            {
              rotation: quat.create(),
              translation: vec3.fromValues(0.5, 0, -0.5),
              scale: vec3.fromValues(0.3, 0.3, 0.3),
              color: vec4.fromValues(1, 1, 1, 1),
            },
          ],
          buffers: planeBuffers,
          pipeline: pipelineWithTexture,
          withTexture: true,
        },
        {
          instances: [
            {
              rotation: quat.create(),
              translation: vec3.fromValues(-0.5, 0, .5),
              scale: vec3.fromValues(0.3, 0.3, 0.3),
              color: vec4.fromValues(0, 1, 0, 1),
            },
            {
              rotation: quat.create(),
              translation: vec3.fromValues(0.5, 0, 0.5),
              scale: vec3.fromValues(0.3, 0.3, 0.3),
              color: vec4.fromValues(0, 0, 1, 1),
            },
          ],
          buffers: triangleBuffers,
          pipeline: pipelineWithoutTexture,
        },
      ];
      objects.forEach((object) => {
        const instanceArrayBuffer = new Float32Array(object.instances.length * (4 * 4 + 4));
        object.instanceBuffer = device.createBuffer({
          size: instanceArrayBuffer.byteLength,
          usage: GPUBufferUsage.VERTEX,
          mappedAtCreation: true,
        });
        object.instances.forEach((instance, i) => {
          const modelMatrix = mat4.fromRotationTranslationScale(mat4.create(), instance.rotation, instance.translation, instance.scale);
          const offset = i * (4 * 4 + 4);
          instanceArrayBuffer.set(modelMatrix, offset);
          instanceArrayBuffer.set(instance.color, offset + 4 * 4);
        });
        new Float32Array(object.instanceBuffer.getMappedRange()).set(instanceArrayBuffer);
        object.instanceBuffer.unmap();

        // Create Scene Buffer Bind Group
        object.sceneBindGroup = device.createBindGroup({
          layout: object.pipeline.getBindGroupLayout(0),
          entries: [
            {
              binding: 0,
              resource: {
                buffer: sceneBuffer,
              },
            },
          ],
        });

        // Create Material Bind Group
        if (object.withTexture) {
          object.materialBindGroup = device.createBindGroup({
            layout: object.pipeline.getBindGroupLayout(1),
            entries: [
              {
                binding: 0,
                resource: device.createSampler(),
              },
              {
                binding: 1,
                resource: imageTexture.createView(),
              },
            ],
          });
        }
      });

      function update() {
        // Update Scene Buffer
        device.defaultQueue.writeBuffer(sceneBuffer, 0, camera.getViewMatrix());
        device.defaultQueue.writeBuffer(sceneBuffer, 4 * 16, camera.getProjectionMatrix());
      }

      function draw(commandEncoder) {
        const passEncoder = commandEncoder.beginRenderPass(passDescriptor);
        objects.forEach((object) => {
          passEncoder.setPipeline(object.pipeline);
          passEncoder.setBindGroup(0, object.sceneBindGroup);
          if (object.materialBindGroup) {
            passEncoder.setBindGroup(1, object.materialBindGroup);
          }
          passEncoder.setVertexBuffer(0, object.buffers.vertexBuffer);
          passEncoder.setVertexBuffer(1, object.instanceBuffer);
          passEncoder.setIndexBuffer(object.buffers.indexBuffer);
          passEncoder.drawIndexed(object.buffers.count, object.instances.length);
        });
        passEncoder.endPass();
      }

      function render() {
        const textureView = swapChain.getCurrentTexture().createView();
        passDescriptor.colorAttachments[0].attachment = outputTextureView;
        passDescriptor.colorAttachments[0].resolveTarget = textureView;
        const commandEncoder = device.createCommandEncoder();
        update();
        draw(commandEncoder);
        device.defaultQueue.submit([commandEncoder.finish()]);
        requestAnimationFrame(render);
      }

      requestAnimationFrame(render);
    })();
  </script>
</body>

</html>